#!/usr/bin/env python
# encoding: utf-8
"""
erk.py

Explicit Runge-Kutta methods in ODEPY.

ODEPY is a Python module which
consists of a library of time stepping integrators,
as well as utility functions. It is meant to facilitate the
construction and analysis of numerical integrators for 
initial value problems.

The module requires a full installation of Numpy/Scipy.

The code is available under GNU GPL v3 licence at
http://code.google.com/p/odepy

Created by Klas Modin on 2011-04-15.
"""
from numpy import array, dot, zeros, diag, ones, triu, shape, hstack
from odepy.core.core import OneStepMethod

class ExplicitRungeKutta(OneStepMethod):
	"""
	Class for explicit Runge-Kutta methods.
	
	Attributes
	----------
	A : ndarray, shape (s,s)
		The A matrix for the method, where s is the number of stages.
		Matrix should be strictly lower triangular (so that the method is explicit).
	
	b : ndarray, shape (s,)
		The b vector for the method.
	
	f : function
		The RHS of the equation. Should have the form
			"def f(y)"
		and should return an ndarray of same shape as y.
		The input y should not be overwritten by f.
	
	k : ndarray, shape (s,:)
		Array for the stage values k_i.
		(Available after at least one step is taken.)
	
	Methods
	-------
	integrate
		Carry out the integration process.
	"""
	
	name = "ERK"
	
	def __init__(self, f, A, b):
		"""
		Constructor.
		
		Parameters
		----------
		f : function
			The RHS of the equation. Should have the form
				"def f(y)"
			and should return an ndarray of same shape as y.
		
		A : ndarray, shape (s,s)
			The A matrix in the Butcher tableau for the method, 
			where s is the number of stages.
			Matrix should be strictly lower triangular 
			(so that the method is explicit).
		
		b : ndarray, shape (s,)
			The b vector in the Butcher tableau for the method.
		"""
		super(ExplicitRungeKutta, self).__init__()
		dims = A.shape
		if (dims[0] != 1 and dims[0] != dims[1]):
			raise AttributeError("A must be a square matrix.")
		s = dims[0]
		if (s != 1 and (triu(A)!=0).sum()):
			raise AttributeError("The matrix A must be strictly lower triangular.")
		self.A = array(A)
		if (len(b) != s):
			raise AttributeError("Size of matrix A and vector b does not match.")
		self.b = array(b)
		self.f = f
	
	def __call__(self,h,y):
		s = len(self.b)
		if (hasattr(self,'k') == False):
			self.k = zeros(hstack((s,shape(y))),dtype=y.dtype)
		self.k[0] = self.f(y)
		for i in xrange(1,s):
			self.ytmp = y + h*dot(self.A[i,:i],self.k[:i])
			self.k[i] = self.f(self.ytmp)
		y += h*dot(self.b,self.k)
		return y
	


# Shortcuts and efficient implementations of specific ERK methods.
class ForwardEuler(ExplicitRungeKutta):
	"""docstring for ForwardEuler"""
	name = "Euler"
	def __init__(self, f,*args,**kwargs):
		super(ForwardEuler, self).__init__(f=f,A=zeros(1),b=ones(1),*args,**kwargs)
	
	def __call__(self,h,y):
		# More efficient implementation than general ERK.
		y += h*self.f(y)
		return y
	

class Heun(ExplicitRungeKutta):
	"""docstring for Heun"""
	name = "Heun"
	def __init__(self,f,*args,**kwargs):
		super(Heun, self).__init__(f=f,A=diag([1],-1),b=ones(2)/2.0,*args,**kwargs)
	

class ExplicitMidpoint(ExplicitRungeKutta):
	"""docstring for ExplicitMidpoint"""
	name = "ExplicitMidpoint"
	def __init__(self, f,*args,**kwargs):
		super(ExplicitMidpoint, self).__init__(f=f,A=diag([0.5],-1),b=array([0.0,1.0]),*args,**kwargs)
	

class ClassicalRK4(ExplicitRungeKutta):
	"""docstring for ClassicalRK4"""
	name = "RK4"
	def __init__(self,f,*args,**kwargs):
		super(ClassicalRK4, self).__init__(f=f,A=diag([0.5,0.5,1],-1),\
			b=1.0/array([6.0,3.0,3.0,6.0]),*args,**kwargs)
	

if __name__ == '__main__' and __file__.split('/')[-1] == 'erk.py':
	from numpy import hstack, sin, pi, cos
	# from pylab import plot,semilogy,clf,legend
	
	# Setup non-linear pendulum problem
	y0 = hstack((0.99*pi,0))
	def f(y):
		dim = len(y)/2
		return hstack((zeros(dim),-sin(y[:dim])))
	
	def g(y):
		dim = len(y)/2
		return hstack((y[dim:],zeros(dim)))
	
	def fplusg(y):
		dim = len(y)/2
		return hstack((y[dim:],-sin(y[:dim])))
	
	def H(y):
		q = y[:,0]
		p = y[:,1]
		return p*p/2-cos(q)
	
	H0 = H(array([y0]))
	
	solverlist = []
	
	# Euler
	# solverlist.append(ForwardEuler(fplusg))
	
	# Störmer-Verlet
	eulerf = ForwardEuler(f)
	eulerg = ForwardEuler(g)
	solverlist.append(eulerf**0.5*eulerg*eulerf**0.5)
	
	# Heun
	solverlist.append(Heun(fplusg))
	
	# Classical Runge-Kutta
	solverlist.append(ClassicalRK4(fplusg))
	
	# Carry out integration
	tend = 500.0
	stepsize = 0.3324235
	Houtlist = []
	for solver in solverlist:
		(tout,yout) = solver.integrate(tend,y0,stepsize)
		Houtlist.append(abs(H(yout)-H0))
	Hout = array(Houtlist).T
	
	# Plot results
	# clf();semilogy(tout,Hout)
	# legend([str(s) for s in solverlist],loc='best')


